# 我们有一个由平面上的点组成的列表 points。需要从中找出 K 个距离原点 (0, 0) 最近的点。
#  （这里，平面上两点之间的距离是欧几里德距离。）
#  你可以按任何顺序返回答案。除了点坐标的顺序之外，答案确保是唯一的。
#
#  示例 1：
#  输入：points = [[1,3],[-2,2]], K = 1
# 输出：[[-2,2]]
# 解释：
# (1, 3) 和原点之间的距离为 sqrt(10)，
# (-2, 2) 和原点之间的距离为 sqrt(8)，
# 由于 sqrt(8) < sqrt(10)，(-2, 2) 离原点更近。
# 我们只需要距离原点最近的 K = 1 个点，所以答案就是 [[-2,2]]。
#
#  示例 2：
#  输入：points = [[3,3],[5,-1],[-2,4]], K = 2
# 输出：[[3,3],[-2,4]]
# （答案 [[-2,4],[3,3]] 也会被接受。）
import heapq
import math
import random
from typing import List


class Solution:
    def kClosest2(self, points: List[List[int]], k: int) -> List[List[int]]:
        """
        快排
        由于快排每次将小于轴点元素的元素都放到轴点的左边，大于轴点的放到右边，即每次可以确定一个元素的最终位置
        所以只需要某次轴点元素的索引为k,然后返回k之前的所有元素即可
        :param points:
        :param k:
        :return:
        """
        def quickSelect(left: int, right: int) -> None:
            curTargetIndex = random.randint(left, right)
            curTargetDist = points[curTargetIndex][0] ** 2 + points[curTargetIndex][1] ** 2
            points[right], points[curTargetIndex] = points[curTargetIndex], points[right]
            pivotIndex = left - 1
            for j in range(left, right):  # 将比轴点元素小的放到左边，比轴点大的放到右边
                if points[j][0] ** 2 + points[j][1] ** 2 <= curTargetDist:
                    pivotIndex += 1
                    points[pivotIndex], points[j] = points[j], points[pivotIndex]
            pivotIndex += 1  # 最终的轴点位置
            points[pivotIndex], points[right] = points[right], points[pivotIndex]  # 轴点元素归位到该元素最终位置上
            if k - 1 < pivotIndex:
                quickSelect(left, pivotIndex - 1)
            elif k - 1 > pivotIndex:
                quickSelect(pivotIndex + 1, right)

        quickSelect(0, len(points) - 1)
        return points[0:k]

    def kClosest1(self, points: List[List[int]], k: int) -> List[List[int]]:
        """
        堆
        :param points:
        :param k:
        :return:
        """
        heap = []
        for x, y in points:
            dist = math.sqrt(x * x + y * y)
            if len(heap) < k:
                heapq.heappush(heap, (-dist, x, y))
            else:
                if heap[0][0] < -dist:
                    heapq.heapreplace(heap, (-dist, x, y))
        return [[e[1], e[2]] for e in heap]

    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        return self.kClosest2(points, k)


if __name__ == "__main__":
    points = [[3, 3], [5, -1], [-2, 4]]
    K = 2

    # print(random.randint(1, 0))
    # for i in range(2, 2, 1):
    #     print(i)

    # points = [[1, 3], [-2, 2]]
    # K = 1
    print(Solution().kClosest(points, K))
